package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{BucketEncoder, BucketEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite
/**
  * @author YongChen
  * @date 2018/12/25 10:18
  * @description
  * @email yong_chen@sui.com
  */
class BucketEncoderSuite extends FunSuite {

  val paramStr: String =
        """
     |{
     |	"input_pt": "a",
     |	"output_pt": "b",
     | "hive_table": "",
     | "flow_time": "",
     |	"inputCol": "a",
     |	"outputCol": "scaledFeatures",
     | "preserveCols": "",
     |	"modelPath": "" ,
     |  "splits" :[1.0, 3.0, 6.0] ,
     |  "handleInvalid" : "keep"
     | }
    """.stripMargin

  val be = new BucketEncoder()
  val param: BucketEncoderParam = be.parseParam(new ToolClass().encrypt(paramStr))

  test("Bucketizer parameter construction test") {
    assert(
          "a".equals(param.input_pt)
        & "b".equals(param.output_pt)
        & "".equals(param.hive_table)
        & "".equals(param.flow_time)
        & "a".equals(param.inputCol)
        & "scaledFeatures".equals(param.outputCol)
        & "".equals(param.preserveCols)
        & "".equals(param.modelPath)
        & "keep".equals(param.handleInvalid)
    )
  }

  test("Bucketizer training function test") {
    val testingData = TestingDataGenerator.makeImputerData().select("a").filter(!_.getAs[Double](0).isNaN)
    val res = be.train(param, testingData).transform(testingData)
    assert(res.distinct().count() == 2)
  }
}
